// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: tianqiang

#include "coin2/exchange/lbank/api_util/rate_tracker.h"

namespace coin2::exchange::lbank::api_util {

constexpr auto EXAMPLE_SPOT_LBANK = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke18/lbank/trade_key.json",
  "rate_limit_info": {
    "spot_lbank": {
      "place_order": {
        "secs": 10,
        "reqs": 500
      },
      "cancel_order": {
        "secs": 10,
        "reqs": 500
      },
      "query_order": {
        "secs": 10,
        "reqs": 200
      }
    }
  }
},)";

template <class RateTracker>
void CheckFail(RateTracker&& rate_tracker, bool check_result) {
  if (!check_result) {
    LOG(INFO) << "Hit rate limit on [" << rate_tracker->GetName() << "]";
  }
}

LbankRateTracker::LbankRateTracker(
    const MarketExchangeApi& mea,
    const coin2::exchange::order::OrderConnectionConfig& cfg) {
  if (mea.market == coin::proto::MarketType::Spot) {
    RateLimitInfoSpotLbank rate_config;
    RateConfig::Get()->GetConfig(mea, cfg, &rate_config);
    InitRateTracker(&rate_config, mea.String(), EXAMPLE_SPOT_LBANK);
  } else {
    CHECK(false) << "unsupported mea: " << mea.String();
  }
}

void LbankRateTracker::InitRateTracker(
    RateLimitInfoSpotLbank* rate_config,
    std::string_view mea_str,
    std::string_view example) {
  if (!rate_config || !rate_config->has_place_order() || !rate_config->has_cancel_order() ||
      !rate_config->has_query_order()) {
    coin2::exchange::base::order::RateConfig::CheckFail(mea_str, example);
  }

  const auto& place_order = rate_config->place_order();
  const auto& cancel_order = rate_config->cancel_order();
  query_order_rate_limit_info_ = rate_config->query_order();
  LOG(INFO) << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
            << "\nOG request rate was set as below:\n"
            << "\nplace_order: [" << place_order.reqs() << "(reqs)/" << place_order.secs()
            << "(secs)]"
            << "\ncancel_order: [" << cancel_order.reqs() << "(reqs)/" << cancel_order.secs()
            << "(secs)]"
            << "\nquery_order: [" << query_order_rate_limit_info_.reqs() << "(reqs)/"
            << query_order_rate_limit_info_.secs() << "(secs)]"
            << "\n\nOr you can configure customized rate in driver config. For example: " << example
            << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";

  rate_tracker_placeorder_ =
      std::make_unique<FixedWindowRateTracker>(place_order.secs() * 1'000'000'000LL);
  rate_tracker_placeorder_->SetRateLimitInfo(place_order.reqs(), place_order.reqs());
  rate_tracker_placeorder_->SetName("lbank_rate_tracker_placeorder(10 secs)");

  rate_tracker_cancelorder_ =
      std::make_unique<FixedWindowRateTracker>(cancel_order.secs() * 1'000'000'000LL);
  rate_tracker_cancelorder_->SetRateLimitInfo(cancel_order.reqs(), cancel_order.reqs());
  rate_tracker_cancelorder_->SetName("lbank_rate_tracker_cancelorder(10 secs)");

  rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
  rate_limit_report_writer_->SetRateLimitInfoString(rate_config);
}

bool LbankRateTracker::HasEnoughRateLimit(
    int64_t timestamp,
    ActionType action_type,
    const std::string& endpoint) {
  if (CheckEnableTime(timestamp, action_type, endpoint) &&
      CheckRateLimit(timestamp, action_type, endpoint)) {
    return true;
  }

  rate_limit_report_writer_->SetHitRateLimit(timestamp, action_type);
  return false;
}

void LbankRateTracker::Push(
    int32_t num_reqs,
    int64_t timestamp,
    ActionType action_type,
    const std::string& endpoint) {
  if (action_type == ActionType::PLACE) {
    if (rate_tracker_placeorder_) {
      rate_tracker_placeorder_->Push(num_reqs, timestamp);
    }
  } else if (action_type == ActionType::CANCEL) {
    if (rate_tracker_cancelorder_) {
      rate_tracker_cancelorder_->Push(num_reqs, timestamp);
    }
  } else if (action_type == ActionType::QUERY) {
    if (rate_tracker_map_.count(endpoint) == 0) {
      return;
    }
    rate_tracker_map_[endpoint]->Push(num_reqs, timestamp);
  } else {
    return;
  }
}

void LbankRateTracker::SetRateLimitInfoString(const google::protobuf::Message* message) {
  if (!rate_limit_report_writer_) {
    return;
  }
  rate_limit_report_writer_->SetRateLimitInfoString(message);
}

void LbankRateTracker::WriteRateLimitReportMessage(OrderLoggerCB callback) {
  if (!callback) {
    return;
  }
  rate_limit_report_writer_->Write(callback);
}

double LbankRateTracker::GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) {
  switch (action_type) {
    case ActionType::PLACE: {
      return GetMaxRateUsedRatio(rate_tracker_placeorder_, timestamp);
    }
    case ActionType::CANCEL: {
      return GetMaxRateUsedRatio(rate_tracker_cancelorder_, timestamp);
    }
    case ActionType::QUERY: {
      std::vector<double> rate_used_ratios;

      std::transform(
          rate_tracker_map_.begin(),
          rate_tracker_map_.end(),
          std::back_inserter(rate_used_ratios),
          [timestamp, this](const auto& p) { return GetMaxRateUsedRatio(p.second, timestamp); });

      return *max_element(rate_used_ratios.begin(), rate_used_ratios.end());
    }
    default: {
      // shouldn't reached here
      LOG(ERROR) << "ops~ something is wrong, action_type=" << action_type;
      return 0.;
    }
  }
}

bool LbankRateTracker::CheckEnableTime(
    int64_t timestamp,
    ActionType action_type,
    const std::string& endpoint) {
  if (timestamp >= enable_time_) {
    return true;
  }

  LOG_EVERY_N(WARNING, 50) << "Hit api temporary disabled until " << enable_time_ << ", "
                           << action_type << ", " << endpoint;
  return false;
}

bool LbankRateTracker::CheckRateLimit(
    int64_t timestamp,
    ActionType action_type,
    const std::string& endpoint) {
  bool ret = true;
  if (action_type == ActionType::PLACE) {
    if (rate_tracker_placeorder_) {
      ret = rate_tracker_placeorder_->GetRemaining(timestamp) > 1;
      CheckFail(rate_tracker_placeorder_, ret);
    }
  } else if (action_type == ActionType::CANCEL) {
    if (rate_tracker_cancelorder_) {
      ret = rate_tracker_cancelorder_->GetRemaining(timestamp) > 1;
      CheckFail(rate_tracker_cancelorder_, ret);
    }
  } else if (action_type == ActionType::QUERY) {
    if (rate_tracker_map_.count(endpoint) == 0) {
      auto rate_tracker = std::make_unique<FixedWindowRateTracker>(
          query_order_rate_limit_info_.secs() * 1'000'000'000LL);
      rate_tracker->SetRateLimitInfo(
          query_order_rate_limit_info_.reqs(),
          query_order_rate_limit_info_.reqs());
      rate_tracker->SetName("lbank_rate_tracker_endpoint:" + endpoint + "(10 secs)");
      rate_tracker_map_[endpoint] = std::move(rate_tracker);
    } else {
      auto& rate_tracker = rate_tracker_map_[endpoint];
      ret = rate_tracker->GetRemaining(timestamp) > 1;
      CheckFail(rate_tracker, ret);
    }
  }

  return ret;
}
}  // namespace coin2::exchange::lbank::api_util
